/*
使用 select 实现的 TCP 和 UDP echo 服务端。
**/

#include <define.h>
#include <error.h>
#include <wrapfunc.h>
#include <sockfunc.h>
#include <sysfunc.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <errno.h>
#include <string.h>


void 
sig_chld_waitpid(int signo) {
	int status;
	pid_t pid;
	while ((pid=waitpid(-1, &status, WNOHANG)) > 0)
		error_msg("child process %d terminate", pid);
    (void)signo;
}

int
main() {
	Signal(SIGCHLD, sig_chld_waitpid);

	struct sockaddr_in localaddr;
	memset(&localaddr, 0, sizeof(localaddr));
	localaddr.sin_family = AF_INET;
	localaddr.sin_port = htons(SERV_PORT);
	localaddr.sin_addr.s_addr = htonl(INADDR_ANY);

    int reuse = 1;
	int listenfd = Socket(AF_INET, SOCK_STREAM, 0);
    Setsockopt(listenfd, SOL_SOCKET, SO_REUSEADDR, (void*)&reuse, sizeof(reuse));
	Bind(listenfd, (struct sockaddr*)&localaddr, sizeof(localaddr));
	Listen(listenfd, LISTENQ);

	memset(&localaddr, 0, sizeof(localaddr));
	localaddr.sin_family = AF_INET;
	localaddr.sin_port = htons(SERV_PORT);
	localaddr.sin_addr.s_addr = htonl(INADDR_ANY);

	int datagram_fd = Socket(AF_INET, SOCK_DGRAM, 0);
    Bind(datagram_fd, (struct sockaddr*)&localaddr, sizeof(localaddr));
    error_msg("wait for tcp and udp clients");

	struct sockaddr_in connaddr;
	socklen_t connlen;
	char buf[MAX_LINE];
	int maxfd1 = (listenfd>datagram_fd ? listenfd:datagram_fd) + 1;
    int n;
	fd_set rset;
	FD_ZERO(&rset);
	for (;;) {
		FD_SET(listenfd, &rset);
		FD_SET(datagram_fd, &rset);
		n = select(maxfd1, &rset, NULL, NULL, NULL);
        if (n < 0) {
            if (errno == EINTR) {
                error_msg("select is interrupted");
                continue;
            }
            else {
                error_dump("fail to select");
            }
        }

		if (FD_ISSET(listenfd, &rset)) {
			connlen = sizeof(connaddr);
			int connfd = Accept(listenfd, (struct sockaddr*)&connaddr, &connlen);
			error_msg("new connection fd:%d from %s", connfd, sock_ntop((struct sockaddr*)&connaddr, connlen));
			
			if (Fork() == 0) {
				Close(listenfd);
				for (;;) {
					ssize_t n;
					n = Read(connfd, buf, sizeof(buf));
					if (n <= 0) {
                        error_msg("connection fd:%d is closed", connfd);
						Close(connfd);
						return 0;
					}
					Writen(connfd, buf, n);
				}
			}
			Close(connfd);
		}

		if (FD_ISSET(datagram_fd, &rset)) {
			ssize_t n;
			connlen = sizeof(connaddr);
			n = Recvfrom(datagram_fd, buf, sizeof(buf), 0, (struct sockaddr*)&connaddr, &connlen);
			error_msg("new datagram from %s", sock_ntop((struct sockaddr*)&connaddr, connlen));
			Sendto(datagram_fd, buf, n, 0, (struct sockaddr*)&connaddr, connlen);
		}
	}

	Close(listenfd);
	Close(datagram_fd);
	return 0;
}
